# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     metric
   Description :   
   Author :       lth
   date：          2022/2/11
-------------------------------------------------
   Change Activity:
                   2022/2/11 9:14: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import matplotlib.pyplot as plt
import numpy as np
import torch

from config import GetConfig

# 设置显示中文
plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号
'''
used for single box cal
'''


def list_mean(list):
    temp = 0
    for i in list:
        temp += i
    return temp / len(list)


def box_iou(box1, box2, x1y1x2y2=True):
    """
    Returns the IoU of two bounding boxes
    """
    if not x1y1x2y2:
        # Transform from center and width to exact coordinates
        b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
        b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
        b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
        b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
    else:
        # Get the coordinates of bounding boxes
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]

    # get the coordinates of the intersection rectangle
    inter_rect_x1 = max(b1_x1, b2_x1)
    inter_rect_y1 = max(b1_y1, b2_y1)
    inter_rect_x2 = min(b1_x2, b2_x2)
    inter_rect_y2 = min(b1_y2, b2_y2)
    # Intersection area
    inter_area = clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * clamp(
        inter_rect_y2 - inter_rect_y1 + 1, min=0
    )
    # Union Area
    b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
    b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)

    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

    return iou


def clamp(in_put, min=0):
    if in_put > min:
        return in_put
    else:
        return min


def get_classes(class_path="./data/class.txt"):
    with open(class_path, 'r') as f:
        class_names = f.readlines()
    return [c.strip('\n') for c in class_names]


class Map:
    def __init__(self, class_path=None):
        self.ap_iou = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
        # self.cls = get_classes(class_path)
        self.num_class = 1
        self.ap_class = [[0 for _ in range(self.num_class)] for _ in range(len(self.ap_iou))]
        self.cal = [[[] for _ in range(self.num_class)] for _ in range(len(self.ap_iou))]
        self.class_total = [0 + 1e-10 for _ in range(self.num_class)]

    def update(self, pred, target):
        self.get_batch_statistics(pred, target, self.ap_iou)

    def get(self):
        self.cal_precision_recall_ap()

        map = []
        for i in range(len(self.ap_iou)):
            # print(self.ap_class[i], list_mean(self.ap_class[i]))
            map.append(list_mean(self.ap_class[i]))
        plt.bar(x=["map" + str(x) for x in self.ap_iou], height=map)
        plt.title("MAP")
        plt.savefig("data/" + GetConfig.target_dir + "/result_map.jpg")
        plt.close()
        # return self.ap_class, list_mean(self.ap_class[i]) for i in range(len(self.ap_iou))

    def reset(self):
        self.ap_class = [[0 for _ in range(self.num_class)] for _ in range(len(self.ap_iou))]
        self.cal = [[[] for _ in range(self.num_class)] for _ in range(len(self.ap_iou))]
        self.class_total = [0 + 1e-10 for _ in range(self.num_class)]

    '''
    这里输入的是outputs 和 targets ，数据格式为ndarray
    '''

    def cal_precision_recall_ap(self):
        figure, axes = plt.subplots(self.num_class, len(self.ap_iou),
                                    figsize=(20 * len(self.ap_iou), 9 * self.num_class), dpi=100)
        for index, _ in enumerate(self.ap_iou):
            for i in range(self.num_class):
                temp = self.cal[index][i]
                if len(temp) == 0:
                    continue
                temp = np.array(temp)
                temp = temp[np.argsort(-temp[:, 0]), :]
                pre = np.zeros((temp.shape[0], 1))
                rec = np.zeros((temp.shape[0], 1))
                total_pre = 0
                total_rec = 0
                for j in range(temp.shape[0]):
                    total_pre += temp[j, 1]
                    pre[j] = total_pre / (j + 1)
                for j in range(temp.shape[0]):
                    total_rec += temp[j, 1]
                    rec[j] = total_rec / self.class_total[i]

                self.ap_class[index][i], p, r = Map.voc_ap(rec.T[0], pre.T[0])
                axes[i][index].step(r, p)
                axes[i][index].set_title(
                    self.cls[i] + "  AP" + str(self.ap_iou[index]) + ": " + str(self.ap_class[index][i]))
        plt.savefig("data/" + GetConfig.target_dir + "/ap_result.jpg")
        plt.close()

    def get_batch_statistics(self, outputs, targets, iou_threshold):

        # 这里不用了，因为在前面的datalist中，type为test，不做box的scale变化
        # 但是output的size是需要变回去的
        # targets = np.array(targets)
        # targets[:, :4] = targets[:, :4] * 416

        # 第一轮循环 去计算 [score,correct]
        outputs = outputs[torch.argsort(outputs[:, -1], descending=True)]

        for iou_index, iou in enumerate(iou_threshold):
            # 这里的数据格式是 xyxy
            appendix = []
            for sample_i in range(len(outputs)):
                output = outputs[sample_i]
                pred_box = output[:4]
                pred_score = output[-1]
                pred_label = output[-2]
                ans = None

                for index, target in enumerate(targets):
                    # 目标框在里面
                    if box_iou(pred_box, target[:4]) >= iou:
                        if index not in appendix:
                            if pred_label == target[4]:
                                ans = [pred_score, True]
                                appendix.append(index)
                            else:
                                ans = [pred_score, False]
                            break

                # 如果没有预测到 gt_label==0，也就是 ans 没有变化的情况下，用于计算precision
                if ans is None:
                    ans = [pred_score, False]
                # 根据pred_label去append
                self.cal[iou_index][int(pred_label)].append(np.array(ans))
        # 第二轮循环 去计算 各个target的数量，用于计算recall
        for target in targets:
            self.class_total[int(target[4])] += 1

    @staticmethod
    def voc_ap(rec, prec):  # rec:召回率 prec:准确率；召回率越高，准确率越低

        # correct AP calculation
        # first append sentinel values at the end
        mrec = np.concatenate(([0.], rec, [1.]))
        mpre = np.concatenate(([1.], prec, [0.]))

        # compute the precision envelope
        for i in range(mpre.size - 1, 0, -1):
            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

        # to calculate area under PR curve, look for points
        # where X axis (recall) changes value
        i = np.where(mrec[1:] != mrec[:-1])[0]

        # and sum (\Delta recall) * prec
        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])  # 计算PR曲线向下包围的面积
        return ap, mpre, mrec


################################################################################################


class ROC:
    def __init__(self, num_cls):
        self.iou_threshold = [0.5]
        self.num_class = num_cls
        self.cal = [[] for _ in range(len(self.iou_threshold))]
        self.class_total = [0 + 1e-10 for _ in range(self.num_class)]

    def update(self, pred, target):
        self.get_batch_statistics(pred, target, self.iou_threshold)

    def get(self):
        self.cal_roc_curve()

    def reset(self):
        self.cal = [[] for _ in range(len(self.iou_threshold))]
        self.class_total = [0 + 1e-10 for _ in range(self.num_class)]

    def get_batch_statistics(self, outputs, targets, iou_threshold):
        # 第一轮循环 去计算 [score,correct]
        outputs = outputs[torch.argsort(outputs[:, -1], descending=True)]

        for iou_index, iou in enumerate(iou_threshold):
            # 这里的数据格式是 xyxy
            appendix = []
            for sample_i in range(len(outputs)):
                output = outputs[sample_i]
                pred_box = output[:4]
                pred_score = output[-1]
                pred_label = output[-2]
                ans = None

                for index, target in enumerate(targets):
                    # 目标框在里面
                    if box_iou(pred_box, target[:4]) >= iou:
                        if index not in appendix:
                            if pred_label == target[4]:
                                ans = [pred_label.item(), pred_score.item(), True]
                                appendix.append(index)
                            else:
                                ans = [pred_label.item(), pred_score.item(), False]
                            break

                # 如果没有预测到 gt_label==0，也就是 ans 没有变化的情况下，用于计算precision
                if ans is None:
                    ans = [pred_label.item(), pred_score.item(), False]
                # 根据pred_label去append
                self.cal[iou_index].append(np.array(ans))
        # 第二轮循环 去计算 各个target的数量，用于计算recall
        for target in targets:
            self.class_total[int(target[4])] += 1

    def cal_roc_curve(self):
        for index, _ in enumerate(self.iou_threshold):
            temp = self.cal[index]
            if len(temp) == 0:
                continue
            temp = np.array(temp)
            temp = temp[np.argsort(-temp[:, 1]), :]
            TP = np.cumsum(temp[:, -1])
            FP = np.cumsum(1 - temp[:, -1])
            TN = temp.shape[0] - np.sum(temp[:, -1]) - FP
            FN = np.sum(temp[:, -1]) - TP

            TPR = TP / (TP + FN)
            FPR = FP / (FP + TN)

            auc = 0
            for i in range(1, len(FPR)):
                auc += (FPR[i] - FPR[i - 1]) * (TPR[i] + TPR[i - 1]) / 2

            plt.step(FPR, TPR)
            plt.plot([0, 1], [0, 1])
            plt.title("ROC CURVE auc:" + str(auc))
            plt.savefig("data/" + GetConfig.target_dir + "/result_roc.jpg")
            plt.close()
